package de.bk.timecalc.core;

import static org.jscience.mathematics.number.Rational.*;

import java.util.HashMap;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Random;

import org.jscience.mathematics.number.Rational;

public class ProbabilisticDistribution<T>
{
  private static final Random RANDOM = new Random();

  private Map<T, Rational> distribution;
  private Rational sumOfProbabilities;

  public ProbabilisticDistribution()
  {
    this.distribution = new HashMap<T, Rational>();
    this.sumOfProbabilities = Rational.ZERO;
  }

  public ProbabilisticDistribution(
      ProbabilisticDistribution<T> otherDistribution)
  {
    this.distribution = new HashMap<T, Rational>(otherDistribution.distribution);
    this.sumOfProbabilities = Rational.ZERO;
    for (Rational probability : this.distribution.values())
    {
      this.sumOfProbabilities = this.sumOfProbabilities.plus(probability);
    }
  }

  public boolean isEmpty()
  {
    return this.distribution.isEmpty();
  }

  public void add(T item, Rational probability)
  {
    this.checkArguments(item, probability);
    if (this.distribution.containsKey(item))
    {
      throw new IllegalStateException("distribution already contains item "
          + item + ".");
    }
    this.distribution.put(item, probability);
    this.sumOfProbabilities = this.sumOfProbabilities.plus(probability);
  }

  public void add(T item, int i)
  {
    this.add(item, new Rational(i, 1));
  }

  public void remove(T item)
  {
    Rational probability = this.getProbability(item);
    this.distribution.remove(item);
    this.sumOfProbabilities = this.sumOfProbabilities.minus(probability);
  }

  public void normalize()
  {
    if (this.isNormalized())
    {
      return;
    }
    // special case, non empty distribution, but all items have probability 0
    if (!this.isEmpty() && this.sumOfProbabilities.compareTo(ZERO) == 0)
    {
      Rational probability = new Rational(1, this.distribution.size());
      for (Map.Entry<T, Rational> entry : this.distribution.entrySet())
      {
        this.distribution.put(entry.getKey(), probability);
      }
    }
    else
    {
      for (Map.Entry<T, Rational> entry : this.distribution.entrySet())
      {
        this.distribution.put(entry.getKey(), entry.getValue().divide(
            this.sumOfProbabilities));
      }
    }
    this.sumOfProbabilities = Rational.ONE;
  }

  public boolean isNormalized()
  {
    if (this.isEmpty())
    {
      return true;
    }
    return this.sumOfProbabilities.compareTo(Rational.ONE) == 0;
  }

  public void changeProbability(T item, Rational probability)
  {
    this.checkArguments(item, probability);
    this.checkContains(item);
    this.sumOfProbabilities = this.sumOfProbabilities.minus(this.distribution
        .get(item));
    this.distribution.put(item, probability);
    this.sumOfProbabilities = this.sumOfProbabilities.plus(probability);
  }

  private void checkArguments(T item, Rational probability)
  {
    if (item == null)
    {
      throw new IllegalArgumentException("item may not be null.");
    }
    if (probability == null)
    {
      throw new IllegalArgumentException("probability may not be null.");
    }
    if (probability.isNegative())
    {
      throw new IllegalArgumentException("probability may not be negative: "
          + probability);
    }
  }

  public boolean contains(T item)
  {
    return this.distribution.containsKey(item);
  }

  public Rational getProbability(T item)
  {
    this.checkContains(item);
    return this.distribution.get(item);
  }

  private void checkContains(T item)
  {
    if (!this.distribution.containsKey(item))
    {
      throw new NoSuchElementException("There is no item " + item
          + " in this distribution.");
    }
  }

  public T select()
  {
    if (this.isEmpty())
    {
      throw new IllegalStateException(
          "Distribution is empty. You can not select from an empty distribution.");
    }
    this.normalize();

    float randomFloat = RANDOM.nextFloat();
    Rational sumOfProbabilitiesSeenSoFar = ZERO;
    for (Map.Entry<T, Rational> entry : this.distribution.entrySet())
    {
      Rational probability = entry.getValue();
      sumOfProbabilitiesSeenSoFar = sumOfProbabilitiesSeenSoFar
          .plus(probability);
      float sumOfProbabilitiesSeenSoFarAsFloat = sumOfProbabilitiesSeenSoFar
          .floatValue();
      if (sumOfProbabilitiesSeenSoFarAsFloat > randomFloat)
      {
        return entry.getKey();
      }
    }

    // randomFloat is alway smaller than 1 and the last value of
    // sumOfProbabilitiesSeenSoFar should always be 1/1, which means the last
    // value of sumOfProbabilitiesSeenSoFarAsFloat is 1, so the loop above
    // should always return a value.
    throw new AssertionError(
        "This was most unexpected: The last value of sumOfProbabilitiesSeenSoFar was ["
            + sumOfProbabilitiesSeenSoFar + "] and randomFloat is ["
            + randomFloat + "].");
  }
}